/*
 * Copyright 1999-2018 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.csp.sentinel.demo.flow;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.util.TimeUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author jialiang.linjl
 */
public class FlowQpsDemo {

    /**
     * 流控规则的 rule key,即给资源点设置的名字
     */
    private static final String KEY = "abc";

    /**
     * 计数器 用于计算1秒钟内通过的请求数量
     */
    private static AtomicInteger pass = new AtomicInteger();

    /**
     * 被限流阻断的请求数量
     */
    private static AtomicInteger block = new AtomicInteger();

    /**
     * 总的请求数量
     */
    private static AtomicInteger total = new AtomicInteger();

    /**
     * 用于标记这个QPS 限流是否应该停止了
     */
    private static volatile boolean stop = false;

    private static final int threadCount = 32;

    /**
     * 最多就允许打印 100秒的数据，线程就会自动做停止
     */
    private static int seconds = 60 + 40;

    /**
     * 基于 QPS 规则来作限流
     *
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {
        // 初始化流控规则
        initFlowQpsRule();

        /**
         * 开启一个 TimerTask 线程，用于输出请求通过、限流等计数日志
         */
        tick();

        /**
         * 模拟流量
         *
         * first make the system run on a very low condition
         */
        simulateTraffic();

        System.out.println("===== begin to do flow control");
        System.out.println("only 20 requests per second can pass");

    }

    /**
     * 初始化流控规则
     */
    private static void initFlowQpsRule() {
        List<FlowRule> rules = new ArrayList<FlowRule>();

        FlowRule rule1 = new FlowRule();
        rule1.setResource(KEY);
        // set limit qps to 20
        rule1.setCount(20);
        rule1.setGrade(RuleConstant.FLOW_GRADE_QPS);

        // 我的这个资源点是属于那个应用系统的
        rule1.setLimitApp("default");
        rules.add(rule1);

        /**
         * 将流控规则加载
         * 这里面会有 给 listeners 赋值
         */
        FlowRuleManager.loadRules(rules);
    }

    /**
     * 模拟流量
     */
    private static void simulateTraffic() {
        /**
         * 启动 32 个线程，循环不断的对一个资源点发起请求
         */
        for (int i = 0; i < threadCount; i++) {
            // 启动线程
            Thread t = new Thread(new RunTask());
            t.setName("simulate-traffic-Task");
            t.start();
        }
    }

    /**
     * 开启一个 TimerTask 线程
     */
    private static void tick() {
        // 老规矩，TimerTask 需要看一下
        Thread timer = new Thread(new TimerTask());
        timer.setName("sentinel-timer-task");
        timer.start();
    }

    static class TimerTask implements Runnable {

        @Override
        public void run() {
            long start = System.currentTimeMillis();
            System.out.println("begin to statistic!!!");

            long oldTotal = 0;
            long oldPass = 0;
            long oldBlock = 0;
            while (!stop) {
                try {
                    TimeUnit.SECONDS.sleep(1);
                } catch (InterruptedException e) {
                }
                // 1秒钟 内总的请求数
                long globalTotal = total.get();
                long oneSecondTotal = globalTotal - oldTotal;
                oldTotal = globalTotal;

                // 1秒钟内 通过的请求数
                long globalPass = pass.get();
                long oneSecondPass = globalPass - oldPass;
                oldPass = globalPass;

                // 1秒钟内 被限流的请求数
                long globalBlock = block.get();
                long oneSecondBlock = globalBlock - oldBlock;
                oldBlock = globalBlock;

                // 日志记录
                System.out.println(seconds + " send qps is: " + oneSecondTotal);
                System.out.println(TimeUtil.currentTimeMillis() + ", total:" + oneSecondTotal
                    + ", pass:" + oneSecondPass
                    + ", block:" + oneSecondBlock);

                // 当100秒到了后，就停止这个线程的运行
                if (seconds-- <= 0) {
                    stop = true;
                }
            }

            // 日志输出
            long cost = System.currentTimeMillis() - start;
            System.out.println("time cost: " + cost + " ms");
            System.out.println("total:" + total.get() + ", pass:" + pass.get()
                + ", block:" + block.get());
            System.exit(0);
        }
    }

    /**
     * 这个线程就比较核心了，是针对指定的资源点发起请求
     */
    static class RunTask implements Runnable {
        @Override
        public void run() {
            while (!stop) {
                Entry entry = null;

                try {
                    /**
                     * 通过 SphU 这个类的静态方法，entry 资源点进入方法，针对指定的资源点发起一个请求， 想要请求指定的一个资源点，按照名称来的
                     * 这个方法比较深，建议打断点进行调试
                     */
                    entry = SphU.entry(KEY);

                    /**
                     * 即 进入一个指定的资源点就等于针对一个资源点发起一个请求，
                     */

                    // token acquired, means pass
                    pass.addAndGet(1);
                } catch (BlockException e1) {
                    block.incrementAndGet();
                } catch (Exception e2) {
                    // biz exception
                } finally {
                    /**
                     * 进入资源点执行完了一段逻辑，或者说执行完一个接口之后，必须从这个资源点里出来
                     */
                    total.incrementAndGet();
                    if (entry != null) {
                        entry.exit();
                    }
                }

                /**
                 * 一个线程每次发起一个请求后，会休眠一个随机时间
                 */
                Random random2 = new Random();
                try {
                    TimeUnit.MILLISECONDS.sleep(random2.nextInt(50));
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        }
    }
}
